package com.hoppinzq;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.hoppinzq.embedBean.Embed;
import org.apache.http.HttpResponse;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.HttpClientBuilder;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class Embedding {

    public static void main(String[] args) throws IOException {
        List<String> strs=new ArrayList<>();
        strs.add("我爱你");
        strs.add("你爱我");
        Embed embed=new Embedding().getEmbed(strs);
        System.err.println(embed);
    }

    public static Embed getEmbed(List<String> inputs){
        try{
            HttpClient httpClient = HttpClientBuilder.create().build();
            HttpPost request = new HttpPost("http://103.143.11.157:8094/v1/embeddings");
            request.addHeader("content-type", "application/json");
            request.addHeader("Authorization", "Bearer sk-");
            JSONObject json=new JSONObject();
            json.put("model","text-embedding-ada-002");
            json.put("input",inputs);
            json.put("user","zhangqi");
            StringEntity params = new StringEntity(json.toString());
            request.setEntity(params);
            HttpResponse response = httpClient.execute(request);
            int statusCode = response.getStatusLine().getStatusCode();
            BufferedReader in = new BufferedReader(new InputStreamReader(response.getEntity().getContent()));
            String inputLine;
            StringBuffer content = new StringBuffer();
            while ((inputLine = in.readLine()) != null) {
                content.append(inputLine);
            }
            in.close();
            if(statusCode==200){
                JSONObject jsonObject= JSONObject.parseObject(content.toString());
                Map<String, Object> map = (Map<String, Object>) jsonObject;
                // 将Java对象的Map类型转换为目标实体类
                Embed embed = JSONObject.parseObject(JSONObject.toJSONString(map), Embed.class);
                return embed;
            }else{
                return null;
            }
        }catch (Exception ex){
            return null;
        }
    }

    public static Embed getEmbed(String input){
        List<String> temp=new ArrayList<>();
        temp.add(input);
        return getEmbed(temp);
    }
}
